import os
import argparse
import torch
import torch.nn as nn
from torchvision import transforms, datasets
from tqdm import tqdm
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import confusion_matrix, classification_report
from colorama import Fore, Style, init
import pandas as pd

from model import AdvancedResNet50, ResNet50WithAttention, ResNet50Model

# 初始化colorama
init(autoreset=True)

# 设置matplotlib字体，解决中文显示问题
def set_matplotlib_font():
    """设置matplotlib字体以支持中文显示"""
    import matplotlib.font_manager as fm
    # 尝试多种常见中文字体
    chinese_fonts = ['SimHei', 'Microsoft YaHei', 'SimSun', 'NSimSun', 'FangSong', 'KaiTi']
    font_found = False
    
    for font_name in chinese_fonts:
        # 检查字体是否存在
        font_path = fm.findfont(fm.FontProperties(family=font_name))
        if os.path.exists(font_path) and not font_path.endswith('DejaVuSans.ttf'):
            plt.rcParams['font.sans-serif'] = [font_name] + plt.rcParams['font.sans-serif']
            plt.rcParams['axes.unicode_minus'] = False  # 正常显示负号
            print(f"使用中文字体: {font_name}")
            font_found = True
            break
    
    if not font_found:
        print("警告: 未找到支持中文的字体，图表中的中文可能无法正确显示")
        # 使用英文标签作为备选方案
        plt.rcParams['font.sans-serif'] = ['Arial'] + plt.rcParams['font.sans-serif']

def get_transforms(image_size=224):
    """获取测试用的数据变换"""
    test_transform = transforms.Compose([
        transforms.Resize((image_size, image_size)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    return test_transform

def load_model(model_path, model_type='advanced', num_classes=10, device='cuda'):
    """加载预训练模型
    
    Args:
        model_path: 模型权重路径
        model_type: 模型类型，可选值：'advanced', 'attention', 'simple'
        num_classes: 类别数量
        device: 设备
        
    Returns:
        model: 加载好权重的模型
    """
    print(f"{Fore.CYAN}正在加载模型: {model_path}{Style.RESET_ALL}")
    
    # 根据模型类型创建相应的模型
    if model_type == 'advanced':
        model = AdvancedResNet50(num_classes=num_classes)
    elif model_type == 'attention':
        model = ResNet50WithAttention(num_classes=num_classes)
    else:
        model = ResNet50Model(num_classes=num_classes)
    
    # 加载模型权重
    model.load_state_dict(torch.load(model_path, map_location=device))
    model = model.to(device)
    model.eval()
    
    print(f"{Fore.GREEN}模型加载成功！{Style.RESET_ALL}")
    return model

def test_model(model, test_loader, device):
    """测试模型性能
    
    Args:
        model: 模型实例
        test_loader: 测试数据加载器
        device: 设备
        
    Returns:
        accuracy: 准确率
        all_preds: 所有预测结果
        all_targets: 所有真实标签
    """
    correct = 0
    total = 0
    all_preds = []
    all_targets = []
    
    with torch.no_grad():
        for inputs, targets in tqdm(test_loader, desc=f"{Fore.BLUE}测试中{Style.RESET_ALL}"):
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = model(inputs)
            _, predicted = outputs.max(1)
            
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()
            
            # 收集预测结果和真实标签
            all_preds.extend(predicted.cpu().numpy())
            all_targets.extend(targets.cpu().numpy())
    
    accuracy = correct / total
    print(f"{Fore.GREEN}测试准确率: {accuracy*100:.2f}%{Style.RESET_ALL}")
    
    return accuracy, all_preds, all_targets

def plot_confusion_matrix(all_targets, all_preds, class_names, save_path=None):
    """绘制混淆矩阵
    
    Args:
        all_targets: 所有真实标签
        all_preds: 所有预测结果
        class_names: 类别名称列表
        save_path: 保存路径，如果为None则不保存
    """
    # 设置字体
    set_matplotlib_font()
    
    # 提取英文类名作为备选
    english_class_names = [name.split('（')[0] if '（' in name else name for name in class_names]
    
    # 计算混淆矩阵
    cm = confusion_matrix(all_targets, all_preds)
    
    # 保存图表（使用英文标签）
    if save_path:
        plt.figure(figsize=(12, 10))
        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
                    xticklabels=english_class_names, yticklabels=english_class_names)
        plt.xlabel('Predicted Label')
        plt.ylabel('True Label')
        plt.title('Confusion Matrix')
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.close()
        print(f"{Fore.GREEN}混淆矩阵已保存到: {save_path}{Style.RESET_ALL}")
    
    # 显示图表（尝试使用中文标签）
    plt.figure(figsize=(12, 10))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
                xticklabels=class_names, yticklabels=class_names)
    plt.xlabel('预测标签')
    plt.ylabel('真实标签')
    plt.title('混淆矩阵')
    plt.show()

def print_classification_report(all_targets, all_preds, class_names):
    """打印分类报告
    
    Args:
        all_targets: 所有真实标签
        all_preds: 所有预测结果
        class_names: 类别名称列表
    """
    # 生成分类报告
    report = classification_report(all_targets, all_preds, target_names=class_names, digits=4)
    
    print(f"\n{Fore.CYAN}分类报告:{Style.RESET_ALL}")
    print(report)
    
    # 将分类报告转换为DataFrame以便更好地展示
    report_dict = classification_report(all_targets, all_preds, target_names=class_names, digits=4, output_dict=True)
    df = pd.DataFrame(report_dict).transpose()
    
    # 保存分类报告
    report_path = os.path.join('results', 'classification_report.csv')
    os.makedirs('results', exist_ok=True)
    df.to_csv(report_path)
    print(f"{Fore.GREEN}分类报告已保存到: {report_path}{Style.RESET_ALL}")

def main():
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='测试爬行动物和两栖动物分类模型')
    parser.add_argument('--model-path', type=str, default='model/best_model_advanced.pth', help='模型路径')
    parser.add_argument('--model-type', type=str, default='advanced', choices=['advanced', 'attention', 'simple'], help='模型类型')
    parser.add_argument('--data-dir', type=str, default='animal_data/test', help='测试数据集路径')
    parser.add_argument('--batch-size', type=int, default=16, help='批次大小')
    parser.add_argument('--image-size', type=int, default=224, help='图像大小')
    args = parser.parse_args()
    
    # 检查GPU可用性
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"{Fore.CYAN}使用设备: {device}{Style.RESET_ALL}")
    
    # 获取数据变换
    test_transform = get_transforms(args.image_size)
    
    # 加载测试数据集
    test_dataset = datasets.ImageFolder(args.data_dir, transform=test_transform)
    test_loader = torch.utils.data.DataLoader(
        test_dataset, batch_size=args.batch_size, shuffle=False, num_workers=4
    )
    
    # 获取类别名称
    class_names = test_dataset.classes
    print(f"{Fore.YELLOW}类别: {class_names}{Style.RESET_ALL}")
    print(f"{Fore.YELLOW}测试样本数: {len(test_dataset)}{Style.RESET_ALL}")
    
    # 加载模型
    model = load_model(args.model_path, args.model_type, len(class_names), device)
    
    # 测试模型
    accuracy, all_preds, all_targets = test_model(model, test_loader, device)
    
    # 创建结果目录
    os.makedirs('results', exist_ok=True)
    
    # 绘制混淆矩阵
    cm_path = os.path.join('results', 'confusion_matrix.png')
    plot_confusion_matrix(all_targets, all_preds, class_names, cm_path)
    
    # 打印分类报告
    print_classification_report(all_targets, all_preds, class_names)
    
    print(f"\n{Fore.GREEN}测试完成！总准确率: {accuracy*100:.2f}%{Style.RESET_ALL}")

if __name__ == '__main__':
    main()